from torchair._ge_concrete_graph.ge_converter.converter_utils import *


@declare_supported([
    Support(F32(3, 4), F32(3, 4), 0, torch.float32),
    Support(F32(3, 4), F32(3, 4), 1, torch.float32),
])
@register_fx_node_ge_converter(torch.ops.aten._log_softmax_backward_data.default)
def conveter_aten__log_softmax_backward_data_default(
    grad_output: Tensor,
    output: Tensor,
    dim: int,
    input_dtype: int,
    meta_outputs: TensorSpec = None,
):
    """NB: aten::_log_softmax_backward_data(Tensor grad_output, Tensor output, int dim, ScalarType input_dtype) -> Tensor"""
    input_ge_type = torch_type_to_ge_type(input_dtype)
    half_to_float = grad_output.dtype != input_ge_type
    if half_to_float:
        if grad_output.dtype != DataType.DT_FLOAT or input_ge_type != DataType.DT_FLOAT16:
            raise NotImplementedError("expected input and grad types to match,",
                                      " or input to be at::Half and grad to be at::Float")
    return ge.LogSoftmaxGrad(grad_output, output, axis=[dim])


@register_fx_node_ge_converter(torch.ops.aten._log_softmax_backward_data.out)
def conveter_aten__log_softmax_backward_data_out(
    grad_output: Tensor,
    output: Tensor,
    dim: int,
    input_dtype: int,
    *,
    out: Tensor = None,
    meta_outputs: TensorSpec = None
):
    """NB: aten::_log_softmax_backward_data.out(Tensor grad_output, Tensor output, int dim, ScalarType input_dtype, *, Tensor(a!) out) -> Tensor(a!)"""
    raise RuntimeError("torch.ops.aten._log_softmax_backward_data.out ge_converter is not supported!")
